import glob
from MnistConvNet import ConvNet
import torch
from PIL import Image
import matplotlib.pyplot as plt
from torchvision import transforms
import os
import math

model = ConvNet()
model.eval()
model.load_state_dict(torch.load('./ckpt/MnistModel.ckpt'))

# 获取指定目录下的所有图片
file_processed = glob.glob(os.path.join('resources/MnistTestCkpt', '*'))
# 向上取整，行数
rows = math.ceil(len(file_processed)/5)
for i, image in enumerate(file_processed):
    img = Image.open(image)
    # 将图片转化成张量
    transform = transforms.ToTensor()
    images = transform(img)
    images = images.reshape(1, 1, 28, 28)
    plt.subplot(rows, 5, i+1)
    # 去掉坐标
    plt.axis('off')
    output = model(images)
    label = output.argmax(dim=1)
    plt.title('{}'.format(int(label)))
    plt.imshow(img)

plt.show()



